

import matplotlib.pyplot as plt

# x轴数据
import numpy as np
import torch
from autograd import grad, elementwise_grad
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
from mpl_toolkits.mplot3d import Axes3D
from numpy import arange, exp, log, sin

import torch.nn.functional as F


def align_new(beta, x):
    return beta * x / ((beta + 1) - x)

def align_new2(x):
    # return pow(align_new(0.7, x), 2.2)
    return pow(align_new(50, x), 0.5)

def align_org(alpha, x):
    return pow(x, alpha)


def plot(plt, x, y, label):
    plt.plot(x, y, label=label)  # 绘制折线图
    return plt
def d(f, x):
    return elementwise_grad(f)(x)


def plot2d(plt):
    x = arange(0, 1, 0.01)
    plt.figure()

    alpha = [5, 6]
    # beta = [0.064, 0.095, 0.128, 0.160, 0.192]
    # beta = [0.064, 0.095, 0.128]
    beta = [0.195, 0.095]

    # for i in alpha:
    #     label = f'y1_alpha={i}'
    #     f = lambda x: align_org(i, x)
    #     plt = plot(plt, x, d(f, x), label=label)
    #
    # for i in beta:
    #     label = f'y2_beta_i={i}'
    #     f = lambda x: align_new(i, x)
    #     plt = plot(plt, x, d(f, x), label=label)

    plt = plot(plt, x, align_new2(x), label='log')
    plt = plot(plt, x, align_org(alpha=0.5, x=x), label='org')

    # plt = plot(plt, x, align_new(beta=0.095,x=x), label='gmm')
    # plt = plot(plt, x, align_org(alpha=6, x=x), label='org')

    # for i in alpha:
    #     label = f'y1_alpha={i}'
    #     f = lambda x: align_org(i, x)
    #     plt = plot(plt, x, d(f, x), label=label)
    #
    # for i in beta:
    #     label = f'y2_beta_i={i}'
    #     plt = plot(plt, x, align_new(i, x), label=label)

    plt.legend(loc='upper left')
    plt.title('Simple Line Plot')  # 设置标题
    plt.xlabel('X-axis')  # 设置x轴标签
    plt.ylabel('Y-axis')  # 设置y轴标签
    plt.show()  # 显示图形



def plot3d(plt):
    fig = plt.figure()

    ax = fig.add_axes(Axes3D(fig))

    # Make data.
    X = np.arange(0, 1, 0.01)
    Y = np.arange(0, 1, 0.01)
    X, Y = np.meshgrid(X, Y)
    # gm = 0.095
    gm = 0.256
    beta = 0.329
    alpha_overlap = 6
    alpha_score = 0.5
    # Z = align_new(beta=beta, x=X) * align_new(beta=gm, x=Y)

    # Z = align_org(alpha=alpha_overlap, x=X) * align_org(alpha=alpha_score, x=Y)
    Z = align_new(beta=gm, x=X) * align_org(alpha=alpha_score, x=Y)

    # Z = align_new(beta=gm, x=X) * align_org(alpha=alpha_score, x=Y)
    # Plot the surface.
    surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm,

                           linewidth=0, antialiased=False)

    # Customize the z axis.

    ax.set_zlim(0, 1)

    ax.zaxis.set_major_locator(LinearLocator(10))

    ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))

    # Add a color bar which maps values to colors.

    fig.colorbar(surf, shrink=0.5, aspect=5)

    plt.show()

if __name__ == '__main__':
    # plot2d(plt)
    plot3d(plt)



